﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using OptixDotNet;
using OptixDotNet.Math;


using float3 = RayDen.Library.Core.Primitives.Vector;
using Vector3 = RayDen.Library.Core.Primitives.Vector;

namespace RayDen.Optix.Engine.Scene.Entities
{

    public class PreethamSkyModel
    {
        public Material ResultMaterial;

        // Represents one entry from table 2 in the paper
        public struct Datum
        {
            public float wavelength;
            public float sun_spectral_radiance;
            public float k_o;
            public float k_wa;

            public Datum(float w, float s, float ko, float kw)
            {
                wavelength = w;
                sun_spectral_radiance = s;
                k_o = ko;
                k_wa = kw;
            }
        };

        private static float[,] cie_table = new float[38, 4] {
                                                                {380f, 0.0002f, 0.0000f, 0.0007f},
                                                                {390f, 0.0024f, 0.0003f, 0.0105f},
                                                                {400f, 0.0191f, 0.0020f, 0.0860f},
                                                                {410f, 0.0847f, 0.0088f, 0.3894f},
                                                                {420f, 0.2045f, 0.0214f, 0.9725f},

                                                                {430f, 0.3147f, 0.0387f, 1.5535f},
                                                                {440f, 0.3837f, 0.0621f, 1.9673f},
                                                                {450f, 0.3707f, 0.0895f, 1.9948f},
                                                                {460f, 0.3023f, 0.1282f, 1.7454f},
                                                                {470f, 0.1956f, 0.1852f, 1.3176f},

                                                                {480f, 0.0805f, 0.2536f, 0.7721f},
                                                                {490f, 0.0162f, 0.3391f, 0.4153f},
                                                                {500f, 0.0038f, 0.4608f, 0.2185f},
                                                                {510f, 0.0375f, 0.6067f, 0.1120f},
                                                                {520f, 0.1177f, 0.7618f, 0.0607f},

                                                                {530f, 0.2365f, 0.8752f, 0.0305f},
                                                                {540f, 0.3768f, 0.9620f, 0.0137f},
                                                                {550f, 0.5298f, 0.9918f, 0.0040f},
                                                                {560f, 0.7052f, 0.9973f, 0.0000f},
                                                                {570f, 0.8787f, 0.9556f, 0.0000f},

                                                                {580f, 1.0142f, 0.8689f, 0.0000f},
                                                                {590f, 1.1185f, 0.7774f, 0.0000f},
                                                                {600f, 1.1240f, 0.6583f, 0.0000f},
                                                                {610f, 1.0305f, 0.5280f, 0.0000f},
                                                                {620f, 0.8563f, 0.3981f, 0.0000f},

                                                                {630f, 0.6475f, 0.2835f, 0.0000f},
                                                                {640f, 0.4316f, 0.1798f, 0.0000f},
                                                                {650f, 0.2683f, 0.1076f, 0.0000f},
                                                                {660f, 0.1526f, 0.0603f, 0.0000f},
                                                                {670f, 0.0813f, 0.0318f, 0.0000f},

                                                                {680f, 0.0409f, 0.0159f, 0.0000f},
                                                                {690f, 0.0199f, 0.0077f, 0.0000f},
                                                                {700f, 0.0096f, 0.0037f, 0.0000f},
                                                                {710f, 0.0046f, 0.0018f, 0.0000f},
                                                                {720f, 0.0022f, 0.0008f, 0.0000f},

                                                                {730f, 0.0010f, 0.0004f, 0.0000f},
                                                                {740f, 0.0005f, 0.0002f, 0.0000f},
                                                                {750f, 0.0003f, 0.0001f, 0.0000f}
                                                            };
        public static Datum[] data = new Datum[38] {
                  new Datum(  0.38f, 1655.9f, -1f, -1f),
                  new Datum(0.39f, 1623.37f, -1f, -1f),
                  new Datum(0.4f, 2112.75f, -1f, -1f),
                  new Datum(0.41f, 2588.82f, -1f, -1f),
                  new Datum(0.42f, 2582.91f, -1f, -1f),
                  new Datum(0.43f, 2423.23f, -1f, -1f),
                  new Datum(0.44f, 2676.05f, -1f, -1f),
                  new Datum(0.45f, 2965.83f, 0.003f, -1f),
                  new Datum(0.46f, 3054.54f, 0.006f, -1f),
                  new Datum(0.47f, 3005.75f, 0.009f, -1f),
                  new Datum(0.48f, 3066.37f, 0.014f, -1f),
                  new Datum(0.49f, 2883.04f, 0.021f, -1f),
                  new Datum(0.5f, 2871.21f, 0.03f, -1f),
                  new Datum(0.51f, 2782.5f, 0.04f, -1f),
                  new Datum(0.52f, 2710.06f, 0.048f, -1f),
                  new Datum(0.53f, 2723.36f, 0.063f, -1f),
                  new Datum(0.54f, 2636.13f, 0.075f, -1f),
                  new Datum(0.55f, 2550.38f, 0.085f, -1f),
                  new Datum(0.56f, 2506.02f, 0.103f, -1f),
                  new Datum(0.57f, 2531.16f, 0.12f, -1f),
                  new Datum(0.58f, 2535.59f, 0.12f, -1f),
                  new Datum(0.59f, 2513.42f, 0.115f, -1f),
                  new Datum(0.6f, 2463.15f, 0.125f, -1f),
                  new Datum(0.61f, 2417.32f, 0.12f, -1f),
                  new Datum(0.62f, 2368.53f, 0.105f, -1f),
                  new Datum(0.63f, 2321.21f, 0.09f, -1f),
                  new Datum(0.64f, 2282.77f, 0.079f, -1f),
                  new Datum(0.65f, 2233.98f, 0.067f, -1f),
                  new Datum(0.66f, 2197.02f, 0.057f, -1f),
                  new Datum(0.67f, 2152.67f, 0.048f, -1f),
                  new Datum(0.68f, 2109.79f, 0.036f, -1f),
                  new Datum(0.69f, 2072.83f, 0.028f, 0.028f),
                  new Datum(0.7f, 2024.04f, 0.023f, 0.023f),
                  new Datum(0.71f, 1987.08f, 0.018f, 0.018f),
                  new Datum(0.72f, 1942.72f, 0.014f, 0.014f),
                  new Datum(0.73f, 1907.24f, 0.011f, 0.011f),
                  new Datum(0.74f, 1862.89f, 0.01f, 0.01f),
                  new Datum(0.75f, 1825.92f, 0.009f, 0.009f)
                };                  // Table2

        public PreethamSkyModel()
        {
            _turbidity = 2.85f;
            _sun_phi = 2.85f;
            _overcast = 0.95f;
            _dirty = true;
            _sun_theta = (float)Math.PI;
            _up = new Vector3(0f, 0f, -1f);
        }


        public void InitProgram(Program skyMiss)
        {
            preprocess();
            skyMiss["overcast"].Set(_overcast);
            skyMiss["sun_direction"].SetFloat3(_sun_dir);
            skyMiss["sun_color"].SetFloat3(_sun_color);
            skyMiss["sky_up"].SetFloat3(this._up);
            skyMiss["inv_divisor_Yxy"].SetFloat3(this._inv_divisor_Yxy);
            skyMiss["c0"].SetFloat3(this._c0);
            skyMiss["c1"].SetFloat3(this._c1);
            skyMiss["c2"].SetFloat3(this._c2);
            skyMiss["c3"].SetFloat3(this._c3);
            skyMiss["c4"].SetFloat3(this._c4);
        }


        public void InitContext(Context skyMiss)
        {
            _dirty = true;
            preprocess();
            skyMiss["overcast"].Set(_overcast);
            skyMiss["sun_direction"].SetFloat3(_sun_dir);
            skyMiss["sun_color"].SetFloat3(_sun_color);
            skyMiss["sky_up"].SetFloat3(this._up);
            skyMiss["inv_divisor_Yxy"].SetFloat3(this._inv_divisor_Yxy);
            skyMiss["c0"].SetFloat3(this._c0);
            skyMiss["c1"].SetFloat3(this._c1);
            skyMiss["c2"].SetFloat3(this._c2);
            skyMiss["c3"].SetFloat3(this._c3);
            skyMiss["c4"].SetFloat3(this._c4);
        }

        public void preprocess()
        {
            if (!_dirty) return;

            _dirty = false;


            _c0 = new Vector3(0.1787f * _turbidity - 1.4630f,
                                -0.0193f * _turbidity - 0.2592f,
                                -0.0167f * _turbidity - 0.2608f);

            _c1 = new Vector3(-0.3554f * _turbidity + 0.4275f,
                                -0.0665f * _turbidity + 0.0008f,
                                -0.0950f * _turbidity + 0.0092f);

            _c2 = new Vector3(-0.0227f * _turbidity + 5.3251f,
                                -0.0004f * _turbidity + 0.2125f,
                                -0.0079f * _turbidity + 0.2102f);

            _c3 = new Vector3(0.1206f * _turbidity - 2.5771f,
                                -0.0641f * _turbidity - 0.8989f,
                                -0.0441f * _turbidity - 1.6537f);

            _c4 = new Vector3(-0.0670f * _turbidity + 0.3703f,
                                -0.0033f * _turbidity + 0.0452f,
                                -0.0109f * _turbidity + 0.0529f);

            float sun_theta_2 = _sun_theta * _sun_theta;
            float sun_theta_3 = sun_theta_2 * _sun_theta;

            float xi = (4.0f / 9.0f - _turbidity / 120.0f) *
                             ((float)(Math.PI) - 2.0f * _sun_theta);

            float3 zenith = new Vector3();
            // Preetham paper is in kilocandellas -- we want candellas
            zenith.X = ((4.0453f * _turbidity - 4.9710f) * (float)Math.Tan(xi) - 0.2155f * _turbidity + 2.4192f) * 1000.0f;
            zenith.Y = _turbidity * _turbidity * (0.00166f * sun_theta_3 - 0.00375f * sun_theta_2 + 0.00209f * _sun_theta) +
                                    _turbidity * (-0.02903f * sun_theta_3 + 0.06377f * sun_theta_2 - 0.03202f * _sun_theta + 0.00394f) +
                                                 (0.11693f * sun_theta_3 - 0.21196f * sun_theta_2 + 0.06052f * _sun_theta + 0.25886f);
            zenith.Z = _turbidity * _turbidity * (0.00275f * sun_theta_3 - 0.00610f * sun_theta_2 + 0.00317f * _sun_theta) +
                                    _turbidity * (-0.04214f * sun_theta_3 + 0.08970f * sun_theta_2 - 0.04153f * _sun_theta + 0.00516f) +
                                                 (0.15346f * sun_theta_3 - 0.26756f * sun_theta_2 + 0.06670f * _sun_theta + 0.26688f);


            float cos_sun_theta = (float)Math.Cos(_sun_theta);

            float3 divisor_Yxy = (_c0 * expf(_c1)).Add(1.0f) * (_c2 * expf(_c3 * _sun_theta).Add(1.0f) + _c4 * cos_sun_theta * cos_sun_theta);

            _inv_divisor_Yxy = zenith.Div(ref divisor_Yxy);

            // 
            // Direct sunlight
            //
            _sun_color = sunColor();

            float sin_sun_theta = (float)Math.Sin(_sun_theta);
            _sun_dir = new Vector3((float)Math.Cos(_sun_phi) * sin_sun_theta,
                                   (float)Math.Sin(_sun_phi) * sin_sun_theta,
                                   (float)Math.Cos(_sun_theta));
            //  optix::Onb onb( _up );
            //onb.inverse_transform( _sun_dir );

        }


        float3 sunColor()
        {
            preprocess();

            // optical mass
            float cos_sun_theta = (float)Math.Cos(_sun_theta);
            float m = 1.0f / (cos_sun_theta + 0.15f * (float)Math.Pow(93.885f - rad2deg(_sun_theta), -1.253f));
            if (float.IsNaN(m))
                m = 1.0f;

            float[] results = new float[38];
            for (int i = 0; i < 38; ++i)
            {
                results[i] = data[i].sun_spectral_radiance * 10000.0f // convert from 1/cm^2 to 1/m^2;
                                                            / 1000.0f; // convert from micrometer to nanometer

                results[i] *= calculateAbsorption(_sun_theta, m, data[i].wavelength, _turbidity, data[i].k_o, data[i].k_wa);
                if (float.IsNaN(results[i]))
                {
                    results[i] = 1.0f;
                }
            }


            float X = 0.0f, Y = 0.0f, Z = 0.0f;
            for (int i = 0; i < 38; ++i)
            {
                X += results[i] * cie_table[i, 1] * 10.0f;
                Y += results[i] * cie_table[i, 2] * 10.0f;
                Z += results[i] * cie_table[i, 3] * 10.0f;
            }
            float3 result = XYZ2rgb(new Vector3(X, Y, Z).Mul(683.0f)) / 1000.0f; // return result in kcd/m^2

            return result;
        }



        private Vector3 expf(Vector3 v)
        {
            return new Vector3((float)Math.Exp(v.X), (float)Math.Exp(v.Y), (float)Math.Exp(v.Z));
        }

        // Calculate absorption for a given wavelength of direct sunlight
        public static float calculateAbsorption(float sun_theta, // Sun angle from zenith
                                          float m,         // Optical mass of atmosphere
                                          float lambda,    // light wavelength
                                          float turbidity, // atmospheric turbidity
                                          float k_o,       // atten coeff for ozone
                                          float k_wa)    // atten coeff for h2o vapor
        {
            float alpha = 1.3f;                             // wavelength exponent
            float beta = 0.04608f * turbidity - 0.04586f;  // turbidity coefficient
            float ell = 0.35f;                            // ozone at NTP (cm)
            float w = 2.0f;                             // precipitable water vapor (cm)

            float rayleigh_air = (float)Math.Exp(-0.008735f * m * (float)Math.Pow(lambda, -4.08f));
            float aerosol = (float)Math.Exp(-beta * m * (float)Math.Pow(lambda, -alpha));
            float ozone = k_o > 0.0f ? (float)Math.Exp(-k_o * ell * m) : 1.0f;
            float water_vapor = k_wa > 0.0f ? (float)Math.Exp(-0.2385f * k_wa * w * m / (float)Math.Pow((1.0f + 20.07f * k_wa * w * m), 0.45f)) : 1.0f;

            return rayleigh_air * aerosol * ozone * water_vapor;

        }

        private static float3 X_F = new float3(3.2410f, -1.5374f, -0.4986f);
        private static float3 Y_F = new float3(-0.9692f, 1.8760f, 0.0416f);
        private static float3 Z_F = new float3(0.0556f, -0.2040f, 1.0570f);

        // Unit conversion helpers
        static float3 XYZ2rgb(float3 xyz)
        {
            float R = float3.Dot(ref xyz, ref X_F);
            float G = float3.Dot(ref xyz, ref Y_F);
            float B = float3.Dot(ref xyz, ref Z_F);
            return new Vector3(R, G, B);

        }

        static float rad2deg(float rads)
        {
            return rads * 180.0f / (float)(Math.PI);
        }

        // Input parameters
        public float _sun_theta;
        public float _sun_phi;
        public float _turbidity;
        public float _overcast;
        public float3 _up;
        public float3 _sun_color;
        public float3 _sun_dir;

        // Precomputation results
        bool _dirty;
        public float3 _c0;
        public float3 _c1;
        public float3 _c2;
        public float3 _c3;
        public float3 _c4;
        public float3 _inv_divisor_Yxy;
    }
}
